Skip to content

feat: add LoadAtomsDataModule for load-atoms datasets - #603

Open
manishpaulish wants to merge 4 commits into
mir-group:developfrom
manishpaulish:feat/load-atoms-datamodule
Open

feat: add LoadAtomsDataModule for load-atoms datasets#603
manishpaulish wants to merge 4 commits into
mir-group:developfrom
manishpaulish:feat/load-atoms-datamodule

Conversation

@manishpaulish

@manishpaulish manishpaulish commented Jun 12, 2026

Copy link
Copy Markdown
Contributor

Adds LoadAtomsDataModule, a Lightning DataModule that wraps the
load-atoms package to provide
easy access to its dataset collection directly from the NequIP training infrastructure.

Any dataset available via load_atoms.load_dataset() (e.g. C-GAP-17, QM9,
rMD17-aspirin) can be used without manual downloading or format conversion.

The implementation:

  • Downloads and caches the dataset via load-atoms in prepare_data()
  • Writes to extxyz format for consumption by the existing ASEDataset pipeline
  • Supports the standard train_val_test_split, subset_len, transforms,
    include_keys, exclude_keys, and key_mapping arguments consistent with
    other DataModules in this codebase

This was suggested by @cw-tan in #600.

@cw-tan

cw-tan commented Jun 13, 2026

Copy link
Copy Markdown
Collaborator

Hi thanks for the contribution. The first thing I'd flag is that we shouldn't make load-atoms a mandatory dependency. Users who want to use the datamodule should install it themselves, and instruction to install them should be in the docstring. The new datamodule should also be exposed in the docs.

@manishpaulish

manishpaulish commented Jun 14, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the feedback. Removed load-atoms from mandatory dependencies and added an install note to the docstring. Also exposed the class in docs/api/datamodule.rst.

@cw-tan

cw-tan commented Jun 14, 2026

Copy link
Copy Markdown
Collaborator

Thanks @manishpaulish . Have you verified that the datamodule works fine when used in a typical traning config for several examples of load-atoms datasets? Also, have you built the docs to check that there are no errors?

@manishpaulish

manishpaulish commented Jun 14, 2026

Copy link
Copy Markdown
Contributor Author

Verified the datamodule works correctly with C-GAP-17 and C-GAP-20U (both download, load, and batch correctly). Docs build succeeds with make html. The only warning is a pre-existing one about torchsim.NequIPTorchSimCalc, unrelated to this PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@manishpaulish @cw-tan